package etxt2db.annotators;

import java.io.Serializable;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import edu.cmu.minorthird.text.BasicTextLabels;
import edu.cmu.minorthird.text.MonotonicTextLabels;
import edu.cmu.minorthird.text.Span;
import edu.cmu.minorthird.text.TextLabels;
import edu.cmu.minorthird.text.learn.SpanFeatureExtractor;
import edu.cmu.minorthird.ui.Recommended;
import edu.umass.cs.mallet.base.fst.MEMM;
import edu.umass.cs.mallet.base.fst.TokenAccuracyEvaluator;
import edu.umass.cs.mallet.base.types.Alphabet;
import edu.umass.cs.mallet.base.types.FeatureVectorSequence;
import edu.umass.cs.mallet.base.types.Instance;
import edu.umass.cs.mallet.base.types.InstanceList;
import edu.umass.cs.mallet.base.types.Sequence;
import etxt2db.api.ClassificationModel;
import etxt2db.corpus.CompleteTextLabels;
import etxt2db.mallet.MalletLoader;


public class MalletMEMMClassificationModel extends ClassificationModel implements Serializable {
	
	/**
	 * 
	 */
	private static final long serialVersionUID = -7995015419584892863L;
	private MEMM memm;
	private SpanFeatureExtractor featureExtractor;
	List<String> attributes;
	
	public MalletMEMMClassificationModel(TextLabels trainingFile, List<String> attributes) throws ParseException{
		this(trainingFile, attributes, new Recommended.TokenFE());
	}
	
	public MalletMEMMClassificationModel(TextLabels labels, List<String> att, SpanFeatureExtractor fe) throws ParseException{
		featureExtractor = fe;
		attributes = att;
		MalletLoader ml = new MalletLoader();
		InstanceList trainingData = ml.load(labels, attributes,featureExtractor);
		
		
		Alphabet targets = trainingData.getTargetAlphabet();
		System.out.print ("State labels:");
		for (int i = 0; i < targets.size(); i++)
			System.out.print (" " + targets.lookupObject(i));
		System.out.println ("");

		// Print out some feature information
		System.out.println ("Number of features = "+trainingData.getDataAlphabet().size());

		/*MultiSegmentationEvaluator eval =
			new MultiSegmentationEvaluator (new String[] {"GB"},
							new String[] {"GI"});*/

		TokenAccuracyEvaluator eval = new TokenAccuracyEvaluator();


		//CRFByGISUpdate crf = new CRFByGISUpdate (p, null);
		memm = new MEMM(trainingData.getDataAlphabet(), trainingData.getTargetAlphabet());
		//crf.addFullyConnectedStatesForLabels();
		//crf.addStatesForLabelsConnectedAsIn (trainingData);
		memm.addStatesForHalfLabelsConnectedAsIn (trainingData);
		memm.setGaussianPriorVariance (Double.POSITIVE_INFINITY);
		//crf.setGaussianPriorVariance (1.0);

		System.out.println("Training on "+trainingData.size()+" training instances");

		//crf.train (trainingData, null, null, eval, 350, new AGIS(4.0));
		memm.train (trainingData, null, null, eval, 100);
	}
	

	public void annotate(MonotonicTextLabels text) {
		MalletLoader ml = new MalletLoader();
		InstanceList testingData;
		try {
			testingData = ml.load(text, attributes,featureExtractor);
		} catch (ParseException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
			return;
		}
		
		Map<String,Sequence> result = new HashMap<String,Sequence>();
	    // Iterate over the sequences in the InstanceList.
	    for (InstanceList.Iterator i = testingData.iterator();i.hasNext();) {
	    	Instance ins = i.nextInstance();
	    	
	        FeatureVectorSequence sequence = (FeatureVectorSequence)ins.getData();
	    
	        // Use the CRF to determine most likely labeling of the 
			// sequence.
	        Sequence labels = memm.viterbiPath(sequence).output();
	        
	        result.put((String)ins.getSource(), labels);
	        
			// print sequences to file
			/*for (int k = 0; k < labels.size(); k++) {
			    System.out.println(labels.get(k));
		    }*/
			//System.out.println();
	    }
	    convertSequenceTextLabels(text, result);
	}

	@Override
	public TextLabels annotatedCopy(TextLabels arg0) {
		//Just create a new TextLabels (copy) before annotate
		CompleteTextLabels copy = new CompleteTextLabels(arg0.getTextBase());
		annotate(copy);
		//Return the new TextLabels
		return copy;
	}

	@Override
	public String explainAnnotation(TextLabels arg0, Span arg1) {
		return null;
	}
	
	private void convertSequenceTextLabels(MonotonicTextLabels text, Map<String,Sequence> labels){
		Iterator<Span> iter = text.getTextBase().documentSpanIterator();
		while(iter.hasNext()){
			Span current = iter.next();
			int i = 0;
			List<Span> list = new ArrayList<Span>();
			while(i<current.size()){
				Span tok = current.subSpan(i, 1);
				if(!tok.asString().matches("\\s*")){
					list.add(tok);
				}
				i++;
			}
			
			Sequence textLabels = labels.get(current.getDocumentId());
			
			int currentSpanBegin=0;
			int currentSpanEnd=0;
			
			
			for(int j=0; j<textLabels.size(); j++){
				if(((String)textLabels.get(j)).endsWith("Unique")){
					String tag = ((String)textLabels.get(j));
					currentSpanBegin = list.get(j).getLoTextToken();
					currentSpanEnd = list.get(j).getLoTextToken()+1;
					Span newSpan=current.charIndexProperSubSpan(currentSpanBegin,currentSpanEnd);
					text.addToType(newSpan,tag.substring(0, tag.lastIndexOf("Unique")));
				}else if(((String)textLabels.get(j)).endsWith("Begin")){
					currentSpanBegin = list.get(j).getLoTextToken();
				}else if(((String)textLabels.get(j)).endsWith("End")){
					String tag = ((String)textLabels.get(j));
					currentSpanEnd = list.get(j).getLoTextToken()+1;
					Span newSpan=current.subSpan(currentSpanBegin, currentSpanEnd-currentSpanBegin);
					text.addToType(newSpan,tag.substring(0, tag.lastIndexOf("End")));
				}
			}
		}
	}
}
